#!/usr/bin/python

import psutil
import re
import socket
import requests
from datetime import datetime
from absl import app, flags

flags.DEFINE_string('process_name_regex', '.*', 'regex match process name.')
flags.DEFINE_integer('expect_match_count', -1, 'expect match count of process name regex.')
flags.DEFINE_string(
  'message_template',
  'process pid={pid}\tname={name}\tcreate_time={create_time}\tcpu_percent={cpu_percent}',
  'slack message template in python "".format() format. support {pid}, {name}, {create_time}, {cpu_percent}.'
)
flags.DEFINE_string('channel', None, 'channel message send to.')

TOKEN = 'xoxb-5035662710-1475501716851-PZKcBOmLPRxHq9SoUTjmJgRn'

kimchi2_proxies = {
   'http': 'http://127.0.0.1:3128',
   'https': 'http://127.0.0.1:3128',
}

def match_process_name(name: str):
  return re.search(flags.FLAGS.process_name_regex, name) != None

def get_all_match_process():
  match_process = []

  for process in psutil.process_iter():
    if match_process_name(process.name()):
      match_process.append(process)

  return match_process

def format_message_template(process: psutil.Process):
  message = flags.FLAGS.message_template

  create_time = datetime \
    .fromtimestamp(process.create_time()) \
    .isoformat(' ', timespec="seconds")

  message = message.format(
    pid=process.pid,
    name=process.name(),
    create_time=create_time,
    cpu_percent=process.cpu_percent(0.1),
  )

  return message + "\n"

def get_slack_message(match_process: list):
  if flags.FLAGS.expect_match_count == len(match_process):
    return None

  message = f"""host:\t{socket.getfqdn()}
time:\t{datetime.now()}
process match regex:\t{flags.FLAGS.process_name_regex}
expect match count:\t{flags.FLAGS.expect_match_count}
actual match count:\t{len(match_process)}
============== match process ==============
"""
  for process in match_process:
    message += format_message_template(process)

  return message

def send_slack_message(message):
  if message == None:
    return

  # print(message)
  send_to_slack(message, flags.FLAGS.channel, 'msg', proxies=kimchi2_proxies)

def get_member_ids(display_name_list, proxies=None):
  url = 'https://slack.com/api/users.list'
  msg = {'token': TOKEN}
  r = requests.post(url, data=msg, proxies=proxies)
  assert r.json()['ok'], ("Failed to get user info ", r.json())
  members_info = r.json()['members']
  name_id_map = {info['name']: info['id'] for info in members_info}
  return [name_id_map.get(name) for name in display_name_list]

def send_to_slack(msg, channel, msg_type, title=None, mention_list=None, proxies=None):
  assert msg_type in ('msg', 'file'), msg_type
  mention_msg = ''
  if mention_list:
    assert isinstance(mention_list, list), f'{mention_list} need to be a list'
    mention_list = [mention.strip() for mention in mention_list]
    trader_member_ids = get_member_ids(mention_list, proxies)
    traders_not_found = ''
    for i, member_id in enumerate(trader_member_ids):
      if member_id is None:
        traders_not_found += '@%s ' % mention_list[i]
      else:
        mention_msg += '<@%s> ' % member_id
    assert len(traders_not_found) == 0, 'Not found in member list: ' + traders_not_found
  if msg_type == 'msg':
    msg = mention_msg + msg
    url = 'https://slack.com/api/chat.postMessage'
    msg = {'token': TOKEN, 'channel': channel, 'text': msg}
  elif msg_type == 'file':
    url = 'https://slack.com/api/files.upload'
    msg = {'token': TOKEN, 'channels': channel, 'content': msg, 'title': title, 'initial_comment': mention_msg}
  r = requests.post(url, data=msg, proxies=proxies)
  assert r.json()['ok'], ("Failed to send ", r.json())

def main(_):
  match_process = get_all_match_process()

  message = get_slack_message(match_process)

  send_slack_message(message)

if __name__ == "__main__":
  app.run(main)
